use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex};

use bytes::Bytes;
use tokio::sync::{broadcast, Notify};
use tokio::time::{self, Duration, Instant};
use tracing::debug;

#[derive(Debug)]
struct Entry {
    id: u64,
    data: Bytes,
    expire_at: Option<Instant>,
}

#[derive(Debug)]
struct State {
    // 数据
    entries: HashMap<String, Entry>,
    // 订阅Sender
    pub_sub: HashMap<String, broadcast::Sender<Bytes>>,
    // 过期ID：BTreeMap具备排序功能
    expirations: BTreeMap<(Instant, u64), String>,
    // 下一个ID
    next_id: u64,
    // 是否停机
    shutdown: bool,
}

impl State {
    /// 下个过期时间
    fn next_expiration(&self) -> Option<Instant> {
        self.expirations.keys().next().map(|key| key.0)
    }
}

#[derive(Debug)]
struct Shared {
    // 锁
    state: Mutex<State>,
    // 后台任务通知
    background_task: Notify,
}

impl Shared {
    /// 清理过期的key
    fn purge_expired_keys(&self) -> Option<Instant> {
        let mut state = self.state.lock().unwrap();
        if state.shutdown {
            return None;
        }

        // lock() 返回 MutexGuard，而不是 &mut State
        // * 解引用
        let state = &mut *state;
        let now = Instant::now();
        while let Some((&(when, id), key)) = state.expirations.iter().next() {
            if when > now {
                // 返回最早的不过期时间
                return Some(when);
            }

            state.entries.remove(key);
            state.expirations.remove(&(when, id));
        }
        None
    }

    /// 是否停机
    fn is_shutdown(&self) -> bool {
        self.state.lock().unwrap().shutdown
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Db {
    /// 浅拷贝
    shared: Arc<Shared>,
}

impl Db {
    pub(crate) fn new() -> Db {
        let shared = Arc::new(Shared {
            state: Mutex::new(State {
                entries: HashMap::new(),
                pub_sub: HashMap::new(),
                expirations: BTreeMap::new(),
                next_id: 0,
                shutdown: false,
            }),
            background_task: Notify::new(),
        });

        // 开启后台任务清理过期key
        tokio::spawn(purge_expired_task(shared.clone()));

        Db {
            shared
        }
    }

    pub(crate) fn get(&self, key: &str) -> Option<Bytes> {
        let state = self.shared.state.lock().unwrap();
        state.entries.get(key).map(|entry| entry.data.clone())
    }

    pub(crate) fn set(&self, key: String, value: Bytes, expire: Option<Duration>) {
        let mut state = self.shared.state.lock().unwrap();

        let id = state.next_id;
        state.next_id += 1;

        // 是否需要提醒后台任务清理
        let mut notify = false;
        // 当前过期时间 < 已存在最早的过期时间 : 应该提醒后台任务清理
        let expire_at = expire.map(|duration| {
            let when = Instant::now() + duration;

            notify = state.next_expiration()
                .map(|expiration| when < expiration)
                .unwrap_or(true);

            //插入过期时间
            state.expirations.insert((when, id), key.clone());

            when
        });

        let prev = state.entries.insert(key, Entry {
            id,
            data: value,
            expire_at,
        });

        if let Some(prev) = prev {
            if let Some(when) = prev.expire_at {
                state.expirations.remove(&(when, prev.id));
            }
        }

        // 在通知后台任务之前释放互斥锁
        // 这有助于减少争用，避免后台任务醒来时由于该函数仍然持有互斥锁而无法获取互斥锁。
        drop(state);

        // 提醒后台清理过期key
        if notify {
            self.shared.background_task.notify_one();
        }
    }

    /// 订阅
    pub(crate) fn subscribe(&self, channel: String) -> broadcast::Receiver<Bytes> {
        use std::collections::hash_map::Entry;

        let mut state = self.shared.state.lock().unwrap();
        match state.pub_sub.entry(channel) {
            // 已经存在channel
            Entry::Occupied(e) => {
                e.get().subscribe()
            }
            // channel 不存在
            Entry::Vacant(e) => {
                let (tx, rx) = broadcast::channel(1024);
                e.insert(tx);
                rx
            }
        }
    }

    /// 发布
    pub(crate) fn publish(&self, channel: &str, msg: Bytes) -> usize {
        let state = self.shared.state.lock().unwrap();

        state.pub_sub.get(channel)
            .map(|tx| tx.send(msg).unwrap_or(0))
            .unwrap_or(0)
    }

    /// 关闭清理任务
    fn shutdown_purge_task(&self) {
        let mut state = self.shared.state.lock().unwrap();
        state.shutdown = true;
        drop(state);
        self.shared.background_task.notify_one();
    }
}

/// 清理过期key
async fn purge_expired_task(shared: Arc<Shared>) {
    while !shared.is_shutdown() {
        if let Some(when) = shared.purge_expired_keys() {
            tokio::select! {
                // 休眠到下次过期时间
                _ = time::sleep_until(when) => {}
                _ = shared.background_task.notified() =>{}
            }
        } else {
            shared.background_task.notified().await;
        }
    }

    debug!("Purge background task shut down")
}

/// Db实例的包装器。在删除时，通过向后台清除任务发出关闭信号，允许有序地清除Db
#[derive(Debug)]
pub(crate) struct DbDropGuard {
    // 删除DbDropGuard会关闭db
    db: Db,
}

impl DbDropGuard {
    pub(crate) fn new() -> DbDropGuard {
        DbDropGuard {
            db: Db::new(),
        }
    }

    pub(crate) fn db(&self) -> Db {
        self.db.clone()
    }
}

impl Drop for DbDropGuard {
    fn drop(&mut self) {
        // drop时关闭清理任务
        self.db.shutdown_purge_task();
    }
}